package com.gitee.wsl.flow.create.state

import androidx.compose.runtime.AbstractApplier
import androidx.compose.runtime.Composable
import androidx.compose.runtime.Composition
import androidx.compose.runtime.Recomposer
import androidx.compose.runtime.snapshots.ObserverHandle
import androidx.compose.runtime.snapshots.Snapshot
import com.gitee.wsl.flow.F
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.CoroutineStart
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.launch
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.coroutines.cancellation.CancellationException

 fun <T> F.stateFlow(body: @Composable () -> T) = channelFlow {
    launchMolecule(
        emitter = {
            trySend(it).getOrThrow()
        },
        body = body,
    )
}

/**
 * Launch a coroutine into this [CoroutineScope] which will continually recompose `body`
 * to produce a [StateFlow] stream of [T] values.
 *
 * The coroutine context is inherited from the [CoroutineScope].
 * Additional context elements can be specified with [context] argument.
 */
fun <T> CoroutineScope.launchMolecule(
    context: CoroutineContext = EmptyCoroutineContext,
    body: @Composable () -> T,
): StateFlow<T> {
    var flow: MutableStateFlow<T>? = null
    launchMolecule(
        context = context,
        emitter = { value ->
            val outputFlow = flow
            if (outputFlow != null) {
                outputFlow.value = value
            } else {
                flow = MutableStateFlow(value)
            }
        },
        body = body,
    )
    return flow!!
}


/**
 * Launch a coroutine into this [CoroutineScope] which will continually recompose `body`
 * in the optional [context] to invoke [emitter] with each returned [T] value.
 *
 * [launchMolecule]'s [emitter] is always free-running and will not respect backpressure.
 * Use [moleculeFlow] to create a backpressure-capable flow.
 *
 * The coroutine context is inherited from the [CoroutineScope].
 * Additional context elements can be specified with [context] argument.
 */
 fun <T> CoroutineScope.launchMolecule(
    emitter: (value: T) -> Unit,
    context: CoroutineContext = EmptyCoroutineContext,
    body: @Composable () -> T,
) {
    /*val clockContext = when (mode) {
        RecompositionMode.ContextClock -> EmptyCoroutineContext
        RecompositionMode.Immediate -> GatedFrameClock(this)
    }*/
    val finalContext = coroutineContext + context //+ clockContext

    val recomposer = Recomposer(finalContext)
    val composition = Composition(UnitApplier, recomposer)
    var snapshotHandle: ObserverHandle? = null
    launch(finalContext, start = CoroutineStart.UNDISPATCHED) {
        try {
            recomposer.runRecomposeAndApplyChanges()
        } catch (e: CancellationException) {
            composition.dispose()
            snapshotHandle?.dispose()
        }
    }

    var applyScheduled = false
    snapshotHandle = Snapshot.registerGlobalWriteObserver {
        if (!applyScheduled) {
            applyScheduled = true
            launch(finalContext) {
                applyScheduled = false
                Snapshot.sendApplyNotifications()
            }
        }
    }

    composition.setContent {
        emitter(body())
    }
}

private object UnitApplier : AbstractApplier<Unit>(Unit) {
    override fun insertBottomUp(index: Int, instance: Unit) {}
    override fun insertTopDown(index: Int, instance: Unit) {}
    override fun move(from: Int, to: Int, count: Int) {}
    override fun remove(index: Int, count: Int) {}
    override fun onClear() {}
}
